#include "raytrace.h"

#include "lambert.h"
//#include "plane.h"
#include "ray.h"
#include "raytrace_script_object.h"
#include "refraction.h"
#include "scoped_mutex_lock.h"
#include "specular.h"
#include "sphere.h"
#include "torus.h"

#include <ppapi/cpp/completion_callback.h>
#include <ppapi/cpp/var.h>

#include <cmath>
#include <cstdio>
#include <cstring>
#include <limits>

namespace {

const uint32_t kAlphaMask = 0xff000000;
const uint32_t kRedMask = 0x00ff0000;
const uint32_t kGreenMask = 0x0000ff00;
const uint32_t kBlueMask = 0x000000ff;

const uint32_t kAlphaShift = 24;
const uint32_t kRedShift = 16;
const uint32_t kGreenShift = 8;
const uint32_t kBlueShift = 0;

#define FROM_ARGB(a, r, g, b) (((r << kRedShift) & kRedMask) |     \
                               ((g << kGreenShift) & kGreenMask) | \
                               ((b << kBlueShift) & kBlueMask) |   \
                               ((a << kAlphaShift) & kAlphaMask))
#define FROM_RGB(r, g, b) FROM_ARGB(0xFF, r, g, b)

#define CLAMP(x, a, b) std::min<float>(std::max<float>(x, a), b)

void FlushCallback(void* data, int32_t result) {
  static_cast<raytrace::Raytrace*>(data)->set_flush_pending(false);
}

inline raytrace::Vector3 GammaCorrect(const raytrace::Vector3& in) {
  static const float kGamma = 1 / 2.2f;
  return raytrace::Vector3(powf(in.x_, kGamma),
                           powf(in.y_, kGamma),
                           powf(in.z_, kGamma));
}

inline raytrace::Vector3 ToneMap(const raytrace::Vector3& in) {
  raytrace::Vector3 clamped(CLAMP(in.x_,0,1), CLAMP(in.y_,0,1), CLAMP(in.z_,0,1));
  const float L = 0.2126f * clamped.x_ + 0.7152f * clamped.y_ + 0.0722f * clamped.z_;
//  const float nL = L / (1 + L);
//  const float scale = nL / L;
  const float scale = 1.0f / (1.0f + L);
  raytrace::Vector3 out = GammaCorrect(raytrace::Vector3::scale(clamped, scale));
  return out;
}
}  // namespace

namespace raytrace {

Raytrace::Raytrace(PP_Instance instance)
    : pp::Instance(instance),
      script_object_(NULL),
      graphics_2d_context_(NULL),
      pixel_buffer_(NULL),
      current_row_(-1),
      percent_complete_(0),
      flush_pending_(false),
      quit_(false),
      num_rays_(1),
      scene_name_("cornell"),
      raytrace_thread_(0) {
  int errno = pthread_mutex_init(&pixel_buffer_mutex_, NULL);
  if (errno != 0)
    printf("ERROR: mutex init failed: %08x\n", errno);

  errno = pthread_mutex_init(&raw_pixel_buffer_mutex_, NULL);
  if (errno != 0)
    printf("ERROR: mutex init failed: %08x\n", errno);
}

Raytrace::~Raytrace() {
  quit_ = true;
  if (raytrace_thread_) {
    pthread_join(raytrace_thread_, NULL);
  }
  DestroyContext();
  delete raw_pixel_buffer_;
  delete pixel_buffer_;
  pthread_mutex_destroy(&raw_pixel_buffer_mutex_);
  pthread_mutex_destroy(&pixel_buffer_mutex_);
  if (script_object_)
    delete script_object_;
}

void Raytrace::DidChangeView(const pp::Rect& position,
                             const pp::Rect& clip) {
  if (position.size().width() == width() &&
      position.size().height() == height()) {
    return;  // Size didn't change, no need to update anything.
  }

  // Create a new device context with the new size.
  DestroyContext();
  CreateContext(position.size());
  // Delete the old pixel buffer and create a new one.
  ScopedMutexLock scoped_mutex(&pixel_buffer_mutex_);
  delete pixel_buffer_;
  pixel_buffer_ = NULL;
  if (graphics_2d_context_ != NULL) {
    printf("Creating pixel buffer size %d x %d\n",
           graphics_2d_context_->size().width(),
           graphics_2d_context_->size().height());
    pixel_buffer_ = new pp::ImageData(this,
                                      PP_IMAGEDATAFORMAT_BGRA_PREMUL,
                                      graphics_2d_context_->size(),
                                      false);
    if (pixel_buffer_ == NULL)
      printf("ERROR: Failed to create pixel buffer %d x %d\n",
             graphics_2d_context_->size().width(),
             graphics_2d_context_->size().height());
  }

  ScopedMutexLock scoped_mutex_raw(&raw_pixel_buffer_mutex_);
  delete[] raw_pixel_buffer_;
  raw_pixel_buffer_ = new Vector3[width() * height()];

  // Fill with a nice bright pink
  for (int i = 0; i < width() * height(); ++i) {
    raw_pixel_buffer_[i] = Vector3(1, 0, 1);
  }

  // Restart the trace if the view size changed while we were rendering.
  if (raytrace_thread_) {
    Start();
  }
}

pp::Var Raytrace::GetInstanceObject() {
  if (!script_object_) {
    script_object_ = new RaytraceScriptObject(this);
    if (script_object_ == NULL)
      printf("ERROR: Failed to create script object\n");
    script_object_->RegisterMethod("paint", &Raytrace::Paint);
    script_object_->RegisterMethod("start", &Raytrace::Start);
    script_object_->RegisterMethod("stop", &Raytrace::Stop);

    script_object_->RegisterProperty("mode", &Raytrace::GetMode,
        &Raytrace::SetMode);
    script_object_->RegisterProperty("scene", &Raytrace::GetScene,
        &Raytrace::SetScene);
    script_object_->RegisterProperty("aa", &Raytrace::GetAa, &Raytrace::SetAa);
    script_object_->RegisterProperty("num_rays", &Raytrace::GetNumRays,
        &Raytrace::SetNumRays);
  }
  return pp::Var(this, script_object_);
}

bool Raytrace::Init(uint32_t argc, const char* argn[], const char* argv[]) {
  return true;
}

void Raytrace::CreateScene() {
  cam_position_ = Vector3(50, 52, 250);
  cam_direction_ = Vector3(0.0f, -0.043f, -1.0f);
  cam_direction_.normalize();

  scene_.Clear();
  
  if (scene_name_ == "cornell") {
    scene_.AddPrimitive(new Sphere(Vector3(10001, 40.8f, 81.6f), 10000,
                            new Lambert(Vector3(0.75f, 0.25f, 0.25f),
                                        Vector3::origin))); // left
    scene_.AddPrimitive(new Sphere(Vector3(-10000 + 99, 40.8f, 81.6f), 10000,
                            new Lambert(Vector3(0.25f, 0.25f, 0.75f),
                                        Vector3::origin))); // right

    scene_.AddPrimitive(new Sphere(Vector3(50, 40.8f, 10000), 10000,
                            new Lambert(Vector3(0.75f, 0.75f, 0.75f),
                                        Vector3::origin))); // back
  //  scene_.AddPrimitive(new Sphere(Vector3(50, 40.8f, -10000 + 250), 10000,
  //                          new Lambert(Vector3::origin,
  //                                      Vector3::origin))); // front

    scene_.AddPrimitive(new Sphere(Vector3(50, 10000, 81.6f), 10000,
                            new Lambert(Vector3(0.75f, 0.75f, 0.75f),
                                        Vector3::origin))); // bottom
    scene_.AddPrimitive(new Sphere(Vector3(50, -10000+81.6f, 81.6f), 10000,
                            new Lambert(Vector3(0.75f, 0.75f, 0.75f),
                                        Vector3::origin))); // top

    scene_.AddPrimitive(new Sphere(Vector3(27, 16.5f, 47), 16.5f,
                            new Specular(Vector3(0.999f, 0.999f, 0.999f),
                                         Vector3::origin))); // mirror
    scene_.AddPrimitive(new Sphere(Vector3(73, 16.5f, 78), 16.5f,
                            new Refraction(1.5f, Vector3(0.999f, 0.999f, 0.999f),
                                           Vector3::origin))); // glass
    scene_.AddPrimitive(new Torus(Vector3(50, 80, 81.6), 5, 20,
                            new Lambert(Vector3::origin,
                                        Vector3(12, 12, 12)))); // light
/*
    scene_.AddPrimitive(new Sphere(Vector3(50, 681.6f-0.27f, 81.6f), 600,
                            new Lambert(Vector3::origin,
                                        Vector3(12.0f, 12.0f, 12.0f)))); // light
*/
  } else if (scene_name_ == "sky") {
    const Vector3 centre(50, 40.8f, -860);
    
    scene_.AddPrimitive(new Sphere(Vector3(3000, 0, 6000), 1600,
                            new Lambert(Vector3::origin,
                                        Vector3::scale(Vector3(1, .9, .8),
                                                       1.2e1*1.56*2)))); // sun
    scene_.AddPrimitive(new Sphere(Vector3(3500, 0, 7000), 1560,
                            new Lambert(Vector3::origin,
                                        Vector3::scale(Vector3(1, .5, .05),
                                                      4.8e1*1.56*2)))); // horizon
    scene_.AddPrimitive(new Sphere(Vector3::add(centre, Vector3(0, 0, -200)), 1e5,
                            new Lambert(Vector3::scale(Vector3(.7, .7, 1), .25f),
                                        Vector3::scale(Vector3(.00063842,
                                                               .02001478,
                                                               .28923243),
                                                       6e-2*8)))); // sky
    scene_.AddPrimitive(new Sphere(Vector3(50, -1e6, 0), 1e6,
                            new Lambert(Vector3(.3,.3,.3),
                                        Vector3::origin))); // ground
    scene_.AddPrimitive(new Sphere(Vector3(50, -110048.5, 0), 1.1e6,
                            new Lambert(Vector3::origin,
                                        Vector3::scale(Vector3(.9,.5,.05),
                                                       4)))); // horizon bright
    scene_.AddPrimitive(new Sphere(Vector3(50, -4e4-30, -3000), 4e4,
                            new Lambert(Vector3(.2,.2,.2),
                                        Vector3::origin))); // mountains

    scene_.AddPrimitive(new Sphere(Vector3(22, 26.5, 42), 26.5,
                            new Specular(Vector3(.596,.596,.596),
                                         Vector3::origin))); // white mirror

    scene_.AddPrimitive(new Sphere(Vector3(75, 13, 82), 13,
                            new Refraction(1.2f, Vector3(.92,.92,.92),
                                           Vector3::origin))); // glass
    scene_.AddPrimitive(new Sphere(Vector3(87, 22, 24), 22,
                            new Refraction(1.5f, Vector3(.417,.417,.417),
                                           Vector3::origin))); // glass2
  } else if (scene_name_ == "overlap") {
    scene_.AddPrimitive(new Sphere(Vector3(50+75, 28, 62), 150,
                            new Refraction(1.5f, Vector3(0.93, 0.83, 0.74),
                                           Vector3::origin)));
    scene_.AddPrimitive(new Sphere(Vector3(50+5, -28, 62), 28,
                            new Lambert(Vector3::origin, Vector3(1,1,1))));
    scene_.AddPrimitive(new Sphere(Vector3(50, 28, 62), 300,
                            new Specular(Vector3(0.93, 0.93, 0.93),
                                         Vector3::origin)));

  } else if (scene_name_ == "island") {
    scene_.AddPrimitive(new Sphere(Vector3(50, 580, -1360), 160,
                            new Lambert(Vector3::origin,
                                        Vector3(200, 200, 200)))); // sun
    scene_.AddPrimitive(new Sphere(Vector3(50, -900, -9980), 160,
                            new Lambert(Vector3::origin,
                                        Vector3(20, 20, 20)))); // horizon
    scene_.AddPrimitive(new Sphere(Vector3(50, -20, -1060), 1e5,
                            new Lambert(Vector3(.4, .4, .4),
                                        Vector3(0.0627, 0.188, 0.569)))); // sky
    scene_.AddPrimitive(new Sphere(Vector3(50, -740, -1060), 800,
                            new Refraction(1.5f,
                                           Vector3(.110, .894, 0.996),
                                           Vector3::origin))); // water
    scene_.AddPrimitive(new Sphere(Vector3(50, -740, -1060), 790,
                            new Lambert(Vector3(.24, .18, .024),
                                        Vector3::origin))); // earth
    scene_.AddPrimitive(new Sphere(Vector3(50, -275, -910), 325,
                            new Lambert(Vector3(.32, .24, .032),
                                        Vector3::origin))); // island
    scene_.AddPrimitive(new Sphere(Vector3(50, -225, -893), 275,
                            new Lambert(Vector3(.015, .225, .015),
                                        Vector3::origin))); // grass
  } else {
    printf("ERROR: Unexpected scene %s\n", scene_name_.c_str());
  }
 
  scene_.DumpToLog();
}

bool Raytrace::GetNumRays(pp::Var& num_rays) {
  num_rays = pp::Var((int32_t) num_rays_);
  return true;
}

bool Raytrace::SetNumRays(const pp::Var& num_rays_var) {
  int32_t num_rays;
  if (num_rays_var.is_number())
    num_rays = (num_rays_var.AsInt());
  else if (num_rays_var.is_string()) {
    // TODO: check error condition
    const char* num_rays_str = num_rays_var.AsString().c_str();
    num_rays = strtoul(num_rays_str, NULL, 10);
  } else
    return false;
    
  num_rays_ = std::max(1L, num_rays);
  printf("set num_rays_ to %lu\n", num_rays_);
  return true;
}

bool Raytrace::GetMode(pp::Var& mode) {
  mode = pp::Var((int32_t) scene_.mode());
  return true;
}

bool Raytrace::SetMode(const pp::Var& mode) {
  if (mode.is_number()) {
    int32_t mode_int = (mode.AsInt());
    scene_.set_mode((Scene::Mode) mode_int);
  } else if (mode.is_string()) {
    // TODO: check error condition
    const std::string mode_str = mode.AsString();
    if (mode_str == "whitted")
      scene_.set_mode(Scene::MODE_WHITTED);
    else if (mode_str == "reverse_path")
      scene_.set_mode(Scene::MODE_REVERSE_PATH);
    else if (mode_str == "debug_normal")
      scene_.set_mode(Scene::MODE_DEBUG_NORMAL);
    else if (mode_str == "debug_path_count")
      scene_.set_mode(Scene::MODE_DEBUG_PATH_COUNT);
    else if (mode_str == "debug_variance")
      scene_.set_mode(Scene::MODE_DEBUG_VARIANCE);
    else
      return false;
  } else
    return false;

  printf("set mode_ to %d\n", scene_.mode());
  return true;
}

bool Raytrace::GetScene(pp::Var& scene) {
  scene = pp::Var(scene_name_);
  return true;
}

bool Raytrace::SetScene(const pp::Var& scene) {
  if (scene.is_string()) {
    // TODO: check error condition
    const std::string scene_str = scene.AsString();
    scene_name_ = scene_str;
  } else
    return false;

  printf("set scene_ to %s\n", scene_name_.c_str());
  return true;
}
bool Raytrace::GetAa(pp::Var& aa) {
  aa = pp::Var((int32_t) aa_);
  return true;
}

bool Raytrace::SetAa(const pp::Var& aa_var) {
  int32_t aa;
  if (aa_var.is_number())
    aa = (aa_var.AsInt());
  else if (aa_var.is_string()) {
    // TODO: check error condition
    const char* aa_str = aa_var.AsString().c_str();
    aa = strtoul(aa_str, NULL, 10);
  } else
    return false;
    
  aa_ = std::max(1L, aa);
  printf("set aa_ to %lu\n", aa_);
  return true;
}

// TODO: signal the thread to stop immediately
pp::Var Raytrace::Stop() {
  int32_t errno = 99;
  if (raytrace_thread_) {
    errno = pthread_join(raytrace_thread_, NULL);
    if (errno != 0)
      printf("ERROR: Failed to join thread %08x\n", (unsigned int) errno);
    raytrace_thread_ = NULL;
  }
  return pp::Var(errno);
}

pp::Var Raytrace::Start() {
  Stop();
  CreateScene();
  int32_t errno = pthread_create(&raytrace_thread_, NULL, DoRaytrace, this);
  if (errno != 0)
    printf("ERROR: Failed to create thread %08x\n", (unsigned int) errno);
  return pp::Var(errno);
}

pp::Var Raytrace::Paint() {
  ScopedMutexLock scoped_mutex(&pixel_buffer_mutex_);
  if (!scoped_mutex.is_valid())
    return pp::Var(0.0f);
  uint32_t* pixel_bits = static_cast<uint32_t*>(pixel_buffer_->data());
  // TODO: only copy what changed?
  // Copy and tone map from raw buffer to pixel buffer
  for (int py = 0; py < height(); ++py) {
    for (int px = 0; px < width(); ++px) {
      int pi = px + py * width();
      Vector3 c = ToneMap(raw_pixel_buffer_[pi]);
      c = Vector3::scale(c, 255);
      pixel_bits[pi] = FROM_RGB((uint32_t) c.x_,
                                (uint32_t) c.y_,
                                (uint32_t) c.z_);
    }
  }
  // Draw scanline
  if (current_row_ > 1) {
    for (int px = 0; px < width(); ++px) {
      int pi = px + (current_row_ - 1) * width();
      pixel_bits[pi] = FROM_RGB(255, 0, 0);
    }
  }

  FlushPixelBuffer();

  if (percent_complete_ == 100.0f) {
    pp::Var exception;
    GetWindowObject().Call("onComplete", &exception);
  }
  return pp::Var(percent_complete_);
}

void Raytrace::CreateContext(const pp::Size& size) {
  if (IsContextValid())
    return;
  graphics_2d_context_ = new pp::Graphics2D(this, size, false);
  if (graphics_2d_context_ == NULL)
    printf("ERROR: Failed to create Graphics2D\n");
  if (!BindGraphics(*graphics_2d_context_))
    printf("ERROR: Failed to bind Graphics2D\n");
}

void Raytrace::DestroyContext() {
  if (!IsContextValid())
    return;
  delete graphics_2d_context_;
  graphics_2d_context_ = NULL;
}

void Raytrace::FlushPixelBuffer() {
  if (!IsContextValid())
    return;
  // Note that the pixel lock is held while the buffer is copied into the
  // device context and then flushed.
  graphics_2d_context_->PaintImageData(*pixel_buffer_, pp::Point());
  if (flush_pending())
    return;
  set_flush_pending(true);
  graphics_2d_context_->Flush(pp::CompletionCallback(&FlushCallback, this));
}

void* Raytrace::DoRaytrace(void* param) {
  Raytrace* raytrace = static_cast<Raytrace*>(param);
  raytrace->DoRaytraceImpl();
  return 0;
}

void Raytrace::DoRaytraceImpl() {
  unsigned int seed = 1;
  srand(seed);

  Vector3 screen_basis_x(width() * 0.75f / height(), 0, 0);
  Vector3 screen_basis_y = Vector3::cross(screen_basis_x,
                                          cam_direction_);
  screen_basis_x.normalize();
  screen_basis_y.normalize();
  screen_basis_y = Vector3::scale(screen_basis_y, 0.75f);
  for (int pyi = 0; pyi < height() && !quit(); ++pyi) {
    current_row_ = height() - pyi - 1;

    for (int pxi = 0; pxi < width() && !quit(); ++pxi) {
      //printf("[%d %d]\n", pxi, pyi);
          
      Vector3 color;

      for (uint32_t aay = 0; aay < aa_; ++aay) {
        for (uint32_t aax = 0; aax < aa_; ++aax) {
          Vector3 radiance;
          // TODO: adaptive sampling
          for (uint32_t ri = 0; ri < num_rays_; ++ri) {
            // Poisson disc sampling
            const float r1 = 2 * drand48();
            const float sxx = r1 < 1 ? sqrt(r1) - 1 : 1 - sqrt(2 - r1);
            const float r2 = 2 * drand48();
            const float syy = r2 < 1 ? sqrt(r2) - 1 : 1 - sqrt(2 - r2);

            const float sdx = (((aax + 0.5f + sxx) / 2 + pxi) / width()) -
                              0.5f;
            const float sdy = (((aay + 0.5f + syy) / 2 + pyi) / height()) -
                              0.5f;
            const Vector3 dx = Vector3::scale(screen_basis_x, sdx);
            const Vector3 dy = Vector3::scale(screen_basis_y, sdy);
            Vector3 ray_d = Vector3::add(Vector3::add(dx, dy),
                                         cam_direction_);
            ray_d.normalize();

            const Vector3 ray_o = cam_position_;

            const Ray ray(ray_o, ray_d);

            //printf("  shooting ray [%.3f %.3f %.3f] -> [%.3f %.3f %.3f]\n",
            //       ray_o.x_, ray_o.y_, ray_o.z_, ray_d.x_, ray_d.y_, ray_d.z_);
            Vector3 r = scene_.Trace(ray);
            radiance = Vector3::add(radiance,
                                    Vector3::scale(r, 1.0f / num_rays_));
          }
          color = Vector3::add(color, Vector3::scale(radiance, 
                                                     1.0f / (aa_ * aa_)));
        }
      }

      const uint32_t pi = width() * current_row_ + pxi;
      //printf("Writing color %.3f %.3f %.3f to pi %lu\n",
      //       color.x_, color.y_, color.z_, pi);
      { 
        ScopedMutexLock scoped_mutex(&raw_pixel_buffer_mutex_);
        raw_pixel_buffer_[pi] = color;
      }  
    }
    percent_complete_ = ((float) pyi/height()) * 100;
  }

  percent_complete_ = 100.0f;
}

}  // namespace 
